from transformers import BertTokenizer, BertModel
import torch
from sklearn.metrics.pairwise import cosine_similarity

from transformers import BertTokenizer, BertModel


# 离线模式下使用本地模型
tokenizer = BertTokenizer.from_pretrained('./local_bert_model')
model = BertModel.from_pretrained('./local_bert_model')


# 示例句子，包含多义词和语义一致性场景
sentence1 = "The bank will provide financial support."
sentence2 = "The children played near the river bank."
sentence3 = "The institution aims to support local communities."

# 生成句子的嵌入
def generate_embeddings(sentence, tokenizer, model):
    inputs = tokenizer(sentence, return_tensors="pt", truncation=True, max_length=128)
    outputs = model(**inputs)  # 修正：使用**inputs而不是'inputs
    embeddings = outputs.last_hidden_state
    return embeddings, inputs['input_ids']

# 获取三个句子的嵌入向量
embeddings1, input_ids1 = generate_embeddings(sentence1, tokenizer, model)
embeddings2, input_ids2 = generate_embeddings(sentence2, tokenizer, model)
embeddings3, input_ids3 = generate_embeddings(sentence3, tokenizer, model)

# 提取目标词 "bank" 和 "support" 的嵌入
def get_word_embedding(word, embeddings, input_ids, tokenizer):
    token_index = torch.where(input_ids[0] == tokenizer.convert_tokens_to_ids(word))[0][0]  # 修正：convert_tokens_to_ids
    return embeddings[0][token_index].detach().numpy()

# 提取 "bank" 和 "support" 的嵌入
bank_embedding1 = get_word_embedding("bank", embeddings1, input_ids1, tokenizer)
bank_embedding2 = get_word_embedding("bank", embeddings2, input_ids2, tokenizer)
support_embedding1 = get_word_embedding("support", embeddings1, input_ids1, tokenizer)
support_embedding3 = get_word_embedding("support", embeddings3, input_ids3, tokenizer)

# 计算嵌入之间的余弦相似度
similarity_bank = cosine_similarity([bank_embedding1], [bank_embedding2])[0][0]
similarity_support = cosine_similarity([support_embedding1], [support_embedding3])[0][0]

print(f"bank在不同上下文中的向量相似性: {similarity_bank:.4f}")  # 修正：bank是字符串
print(f"support在语义一致性场景中的向量相似性: {similarity_support:.4f}")  # 修正：support是字符串

# 打印向量信息
print("\nbank在句子 1 中的嵌入向量 (前5维):", bank_embedding1[:5])  # 修正：bank是字符串
print("\nbank在句子 2 中的嵌入向量 (前5维):", bank_embedding2[:5])  # 修正：bank是字符串
print("\nsupport在句子 1 中的嵌入向量 (前5维):", support_embedding1[:5])  # 修正：support是字符串
print("support在句子 3 中的嵌入向量 (前5维):", support_embedding3[:5])  # 修正：support是字符串